import time

import cv2
from insightface.app import FaceAnalysis
from insightface.app.common import Face

root_path = './'
model_name = 'buffalo_ss'
det_thresh = 0.7
det_size = 450


def test01():
    face_obj = FaceAnalysis(root=root_path, name=model_name, allowed_modules=['detection', 'recognition'],
                            providers=['CANNExecutionProvider', 'CPUExecutionProvider'])
    face_obj.prepare(ctx_id=0, det_thresh=det_thresh, det_size=(det_size, det_size))
    for i in range(19):
        filename = './imgs/' + str(i + 1) + '.jpg'
        start = time.time_ns()
        frame = cv2.imread(filename)
        face_obj.get(frame)
        end = time.time_ns()
        print('识别时间')
        print((end - start) / 1000000000)


def test02():
    stream = 'rtsp://admin:kewen666@192.168.31.22:554/Streaming/Channels/101'
    cap = cv2.VideoCapture(stream)
    face_obj = FaceAnalysis(root=root_path, name=model_name, allowed_modules=['detection', 'recognition'],
                            providers=['CANNExecutionProvider', 'CPUExecutionProvider'])
    face_obj.prepare(ctx_id=0, det_thresh=det_thresh, det_size=(det_size, det_size))
    num = 0
    while cap.isOpened():
        fps = cap.get(cv2.CAP_PROP_FPS)
        res, frame = cap.read()
        if res:
            num = num + 1
            if num > 10000:
                num = 0
            if num % fps == 2:
                start = int(time.time_ns())
                face_obj.get(frame)
                end = time.time_ns()
                print('识别时间')
                print((end - start) / 1000000000)


def test03():
    face_obj = FaceAnalysis(root=root_path, name=model_name, allowed_modules=['detection', 'recognition'],
                            providers=['CANNExecutionProvider', 'CPUExecutionProvider'])
    face_obj.prepare(ctx_id=0, det_thresh=det_thresh, det_size=(det_size, det_size))
    filename = './imgs/' + str(20) + '.jpg'
    start = time.time_ns()
    frame = cv2.imread(filename)
    bboxes, kpss = getDetection(face_obj, frame)
    print('检测时间：')
    print((time.time_ns() - start) / 1000000000)
    rec = getRecognition(bboxes, kpss, face_obj, frame)
    print('识别时间：')
    print((time.time_ns() - start) / 1000000000)




def getDetection(face_obj, frame):
    return face_obj.det_model.detect(frame)

def getRecognition(bboxes, kpss, face_obj, frame):
    if bboxes.shape[0] == 0:
        print('无特征')
        return []
    ret = []
    for i in range(bboxes.shape[0]):
        bbox = bboxes[i, 0:4]
        det_score = bboxes[i, 4]
        kps = None
        if kpss is not None:
            kps = kpss[i]
        face = Face(bbox=bbox, kps=kps, det_score=det_score)
        face_obj.models['recognition'].get(frame, face)
        ret.append(face)
    return ret


if __name__ == '__main__':
    test03()
